package org.zoomdev.zoom.hash;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.LongAdder;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class ProductionConsistentHash {
    // 默认每个物理节点对应160个虚拟节点
    private static final int DEFAULT_VIRTUAL_NODES = 160;
    // 虚拟节点环，使用跳表保证有序且线程安全
    private final ConcurrentSkipListMap<Integer, String> virtualNodes = new ConcurrentSkipListMap<>();
    // 物理节点管理
    private final Set<String> physicalNodes = ConcurrentHashMap.newKeySet();
    // 读写锁控制节点变更
    private final ReentrantReadWriteLock rwLock = new ReentrantReadWriteLock();
    private final ReentrantReadWriteLock.ReadLock readLock = rwLock.readLock();
    private final ReentrantReadWriteLock.WriteLock writeLock = rwLock.writeLock();
    // 性能统计
    private final LongAdder requestCounter = new LongAdder();
    private final Map<String, LongAdder> nodeCounters = new ConcurrentHashMap<>();
    private final int virtualNodesPerShard;

    // 哈希算法配置
    private final HashAlgorithm hashAlgorithm;

    public ProductionConsistentHash() {
        this(DEFAULT_VIRTUAL_NODES, HashAlgorithm.MURMUR3_32);
    }

    public ProductionConsistentHash(int virtualNodesPerShard, HashAlgorithm algorithm) {
        this.virtualNodesPerShard = virtualNodesPerShard;
        this.hashAlgorithm = algorithm;
    }

    /**
     * ​
     * 添加物理节点
     */
    public void addNode(String node) {
        writeLock.lock();
        try {
            if (physicalNodes.add(node)) {
                addVirtualNodes(node);
            }
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * ​
     * 移除物理节点
     */
    public void removeNode(String node) {
        writeLock.lock();
        try {
            if (physicalNodes.remove(node)) {
                removeVirtualNodes(node);
            }
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * ​
     * 获取数据对应的节点
     */
    public String getNode(String key) {
        readLock.lock();
        try {
            if (virtualNodes.isEmpty()) {
                throw new IllegalStateException("No available nodes");
            }

            int hash = hashAlgorithm.hash(key);
            Map.Entry<Integer, String> entry = virtualNodes.ceilingEntry(hash);
            String node = (entry != null) ? entry.getValue() : virtualNodes.firstEntry().getValue();

            // 统计
            requestCounter.increment();
            nodeCounters.computeIfAbsent(node, k -> new LongAdder()).increment();

            return node;
        } finally {
            readLock.unlock();
        }
    }

    /**
     * ​
     * 获取节点分布统计
     */
    public Map<String, Long> getDistribution() {
        Map<String, Long> result = new HashMap<>();
        nodeCounters.forEach((k, v) -> result.put(k, v.longValue()));
        return Collections.unmodifiableMap(result);
    }

    /**
     * ​
     * 获取系统负载统计
     */
    public SystemLoadStats getLoadStats() {
        return new SystemLoadStats(
                requestCounter.longValue(),
                virtualNodes.size(),
                physicalNodes.size()
        );
    }

    private void addVirtualNodes(String node) {
        for (int i = 0; i < virtualNodesPerShard; i++) {
            String vNode = node + "#VN" + i;
            int hash = hashAlgorithm.hash(vNode);
            virtualNodes.put(hash, node);
        }
    }

    private void removeVirtualNodes(String node) {
        virtualNodes.entrySet().removeIf(entry -> entry.getValue().equals(node));
    }


    /**
     * ​
     * 系统负载统计
     */
    public static class SystemLoadStats {
        private final long totalRequests;
        private final int virtualNodeCount;
        private final int physicalNodeCount;

        public SystemLoadStats(long totalRequests, int virtualNodeCount, int physicalNodeCount) {
            this.totalRequests = totalRequests;
            this.virtualNodeCount = virtualNodeCount;
            this.physicalNodeCount = physicalNodeCount;
        }

        // Getters...
    }
}